import { Annotation, AnnotationRoot } from '@langchain/langgraph';
import { BaiLianService, baseInputVars, BL_ImageGenSize, downloadImage, HS_ImageGenSize, LastMsg, LLMNode, LLMProvider } from "../../../core";
import { IEditItem } from '../data';
import { HumanMessage, SystemMessage } from '@langchain/core/messages';

export class EditImgNode extends LLMNode {
    name = "EditImgNode";
    inputVars = Annotation.Root({
        ...baseInputVars.spec,
        image: Annotation<string>,
        sourceImageBase64: Annotation<string>,
        imageBuffer: Annotation<any>,

        IdentifyImgObjectNodeResult: Annotation<any>,
        ImageListExtNodeResult: Annotation<any>,
        editList: Annotation<IEditItem[]>
    });

    outputVars = Annotation.Root({
        lastMsg: Annotation<LastMsg>,
    });

    private async editWithBailian(prompt: string) {
        const svr = new BaiLianService();
        const key = Object.keys(BL_ImageGenSize).find(key => {
            const value = BL_ImageGenSize[key] as string;
            return value.startsWith(this.inputVars.State.imageBuffer.width.toString());
        });
        const size = BL_ImageGenSize[key];
        const res = await svr.GenOrEidtImage(prompt, size, this.inputVars.State.sourceImageBase64);
        return JSON.parse(res.content) as { url: string; }[];
    }

    private async editWithDoubao(prompt: string) {
        const key = Object.keys(BL_ImageGenSize).find(key => {
            const value = BL_ImageGenSize[key] as string;
            return value.startsWith(this.inputVars.State.imageBuffer.width.toString());
        });
        const size = HS_ImageGenSize[key];
        const model = LLMProvider.current.DoubaoSeedream40;
        model.size = size;
        model.prompt = prompt;
        model.image = [
            this.inputVars.State.sourceImageBase64
        ];
        const res = await model.invoke([]);
        return JSON.parse(res.content as any) as { url: string; }[];
    }

    private formatToPrompt(editItem: IEditItem) {
        let editR = `
        ## 当前区域修改要求：
        ${editItem.userInput}

        ** 请保持主体区域外的部分不变 **
        `
        if (editItem.imageExt && editItem.imageExt.length > 0) {
            editR = `
            ## 当前区域修改要求：
            按照图片的概要：${editItem.imageExt} 进行替换，请保持整体的风格统一

            ** 请保持主体区域外的部分不变 **
            `
        }

        let tmp = `
        ## 修改区域ID: ${editItem.objID}
        - 区域坐标<bbox>${editItem.bbox[0].toFixed(2)}, ${editItem.bbox[1].toFixed(2)}, ${editItem.bbox[2].toFixed(2)}, ${editItem.bbox[3].toFixed(2)}</bbox>
        - 区域内的对象信息
        <objectInfo>
        ${editItem.bboxObject}
        </objectInfo>

        ${editR}
        `;

        return tmp;
    }

    private async getImageEditCmd(sourcePrompt: string) {
        const model = LLMProvider.current.ZiJieDeepSeekR1;
        const ouputStream = await model.stream([
            new SystemMessage(`
​角色：​​ 你是一个AI图片编辑助手，专门将用户提供的图片区域修改信息转化为精确的Stable Diffusion inpainting指令。你的输出必须是可直接执行的SD命令，遵循CRAFT原则。

​​输入：​​ 用户提供结构化区域信息，包括区域坐标（bbox）、对象描述和修改要求。

​​输出：​​ 将所有的区域修改整合，生成一个完整的SD inpainting指令，包括正向提示词、负向提示词、重绘参数和质量检查点。指令必须清晰、可执行，并专注于修改指定区域，同时保持整体风格一致。 

** 仅仅只修改区域内容的图像、不要超出区域 **
** 输出的sd 指令不要超过800字 **
** 仅仅输出sd 指令不要输出任何代码 **
            `),
            new HumanMessage(sourcePrompt)
        ]);
        for await (const chunk of ouputStream) { }

        this.usage.add(model.totalUsage);

        return model.lastFullResponse;
    }

    async Run() {
        let prompt = [];
        this.inputVars.State.editList.forEach(n => {
            n.bboxObject = JSON.stringify(this.inputVars.State.IdentifyImgObjectNodeResult.find(x => x.objID == n.objID).bboxObject);
            n.imageExt = this.inputVars.State.ImageListExtNodeResult.find(x => x.objID == n.objID).imageExt;

            prompt.push(this.formatToPrompt(n));
        });

        // const genImageCmd = await this.getImageEditCmd(prompt.join("\n\n"));
        const genImageCmd = prompt.join("\n\n");
        let obj = await this.editWithBailian(genImageCmd);

        for (let item of obj) {
            let filename = await downloadImage(item.url);
            item.url = `http://localhost:9000/coeus_coder/api/image/${filename}`;
        }

        this.outputVars.State.lastMsg = {
            content: JSON.stringify(obj),
            contentType: "image"
        };

        return this.outputVars.State;
    }
}